﻿// Copyright (c) 2014 Morten Bakkedal
// This code is published under the MIT License.

using System;
using System.Collections.Generic;
using System.Text;

using FuncLib.DualFunctions.Collections;
using FuncLib.DualFunctions.Expressions;

namespace FuncLib.DualFunctions.Compilation
{
	internal class FunctionBody
	{
		private DualCodeGenerator generator;
		private string identifier;
		private FunctionParameter[] parameters;
		private List<FunctionBodyStatement> statements;
		private uint nextLocal, nextArray;
		private Dictionary<IndexList, StorageOffsetsReference> offsetsIdentifiers = new Dictionary<IndexList, StorageOffsetsReference>();

		public FunctionBody(DualCodeGenerator generator, string identifier, params FunctionParameter[] parameters)
		{
			this.generator = generator;
			this.identifier = identifier;
			this.parameters = parameters;

			statements = new List<FunctionBodyStatement>();
			nextLocal = 0;
			nextArray = 0;
		}

		public string GetInnerFunctionIdentifier(GeneratedFunction function)
		{
			foreach (FunctionParameter parameter in parameters)
			{
				if (parameter.Function == function)
				{
					return parameter.Identifier;
				}
			}

			throw new IndexOutOfRangeException();
		}

		private void AddStatement(FunctionBodyStatement statement)
		{
			statements.Add(statement);
		}

		public Expression AddLocal(Expression expression)
		{
			string identifier = "c" + nextLocal++;
			AddStatement(new LocalAssignmentStatement(identifier, expression));
			return new LocalReferenceExpression(identifier);
		}

		public void AddFunctionValue(Expression expression)
		{
			AddStatement(new FunctionValueAssignmentStatement(identifier, 0, expression));
		}

		public void AddValues(FunctionValueStorage functionValueStorage, Storage storage, Expression[] expressions)
		{
			if (expressions.Length != storage.Length)
			{
				throw new ArgumentException();
			}

			for (int i = 0; i < expressions.Length; i++)
			{
				AddStatement(new FunctionValueAssignmentStatement(identifier, storage.Index - functionValueStorage.Index + i, expressions[i]));
			}
		}

		public StorageOffsetsReference AddOffsets(IndexList offsets) // do not save as a statement, keep in top wrapped in "unsafe" with other offsets
		{
			StorageOffsetsReference reference;
			if (offsetsIdentifiers.TryGetValue(offsets, out reference))
			{
				// Already used for this function (and loaded on the stack). Return the same reference.
				return reference;
			}

			// Add the offsets to the global list if not already added.
			int index = generator.AddOffsets(offsets);
			//int index = generator.AddIndexList(offsets);

			// Add to the local stack.
			string identifier = "z" + nextArray++;
			AddStatement(new StorageOffsetsAssignmentStatement(identifier, "z[" + index + "]"));
			
			reference = new StorageOffsetsReference(identifier, offsets);
			offsetsIdentifiers.Add(offsets, reference);

			return reference;
		}

		public void AddGradientLoop(FunctionValueStorage functionValueStorage, SortedGradient gradient, Storage gradientStorage, params GradientLoopTerm[] terms)
		{
			// equality comparer with () type cast

			List<GradientLoopTerm> nonZeroTerms = new List<GradientLoopTerm>();
			foreach (GradientLoopTerm term in terms)
			{
				if (!term.Expression.IsZero() && term.Gradient.Count > 0)
				{
					nonZeroTerms.Add(term);
				}
			}
			GradientLoopTerm[][] groupedTerms = QuotientSet<GradientLoopTerm>.Compute(nonZeroTerms.ToArray());

			if (groupedTerms.Length > 1)
			{
				// Add zero assignment statement.
				AddStatement(new LoopZeroStatement(identifier, gradientStorage.Index - functionValueStorage.Index, gradientStorage.Length));
			}

			for (int i = 0; i < groupedTerms.Length; i++)
			{
				// All terms in this group share the same Hessian (but potentially stored with different offsets). Just use the first to generate
				// offsets for the output.
				StorageOffsetsReference z = AddOffsets(groupedTerms[i][0].Gradient.GetIndexList(functionValueStorage.Index, gradient, gradientStorage));

				// Generate offsets for each of the input Hessian.
				foreach (GradientLoopTerm term in groupedTerms[i])
				{
					term.Prepare(this);
				}

				AddStatement(new GradientLoopStatement(identifier, z, groupedTerms[i], groupedTerms.Length > 1));
			}
		}

		public void AddHessianLoop(FunctionValueStorage functionValueStorage, SortedHessian hessian, Storage hessianStorage, params HessianLoopTerm[] terms)
		{
			// equality comparer with () type cast

			List<HessianLoopTerm> nonZeroTerms = new List<HessianLoopTerm>();
			foreach (HessianLoopTerm term in terms)
			{
				if (!term.Expression.IsZero() && term.Hessian.Count > 0)
				{
					nonZeroTerms.Add(term);
				}
			}
			HessianLoopTerm[][] groupedTerms = QuotientSet<HessianLoopTerm>.Compute(nonZeroTerms.ToArray());

			if (groupedTerms.Length > 1)
			{
				// Add zero assignment statement.
				AddStatement(new LoopZeroStatement(identifier, hessianStorage.Index - functionValueStorage.Index, hessianStorage.Length));
			}

			for (int i = 0; i < groupedTerms.Length; i++)
			{
				// All terms in this group share the same Hessian (but potentially stored with different offsets). Just use the first to generate
				// offsets for the output.
				StorageOffsetsReference z = AddOffsets(groupedTerms[i][0].Hessian.GetIndexList(functionValueStorage.Index, hessian, hessianStorage));

				// Generate offsets for each of the input Hessian.
				foreach (HessianLoopTerm term in groupedTerms[i])
				{
					term.Prepare(this);
				}

				AddStatement(new HessianLoopStatement(identifier, z, groupedTerms[i], groupedTerms.Length > 1));
			}
		}

		private void AddLoop(FunctionValueStorage functionValueStorage, Storage loopStorage, LoopTerm[] terms)
		{
			// do as above in generic code
			throw new NotImplementedException();
		}

		public void GenerateCodeNEW(FunctionValueStorage functionValueStorage, out string generatedCode, out string callingParameters)
		{
			callingParameters = functionValueStorage.Index.ToString();
			StringBuilder sb = new StringBuilder();

			sb.Append("(int ");
			sb.Append(identifier);
			foreach (FunctionParameter parameter in parameters)
			{
				sb.Append(", int ");
				sb.Append(parameter.Identifier);
				callingParameters += ", " + parameter.Function.FunctionValueStorage.Index;
			}
			sb.AppendLine(")");

			sb.AppendLine("{");
			foreach (FunctionBodyStatement statement in statements)
			{
				//sb.AppendLine(statement.GenerateCode());
				sb.Append("    " + statement.GenerateCode().Replace("\r\n", "\r\n    ") + "\r\n"); // FIXME remove Replace
			}
			sb.AppendLine("}");

			generatedCode = sb.ToString();
		}

		//public void GenerateCode(out string methodBody, out string callingParameters)

		// GeneratedCode { MethodName, MethodBody, MethodParameters, CallingParameters }

		private class StorageOffsetsAssignmentStatement : FunctionBodyStatement
		{
			private string identifier, assignment;
			//private StorageOffsets offsets;

			public StorageOffsetsAssignmentStatement(string identifier, string assignment)
			{
				this.identifier = identifier;
				this.assignment = assignment;
				//this.offsets = offsets;
			}

			public override string GenerateCode()
			{
				return "int[] " + identifier + " = " + assignment + ";";
				//return "int[] " + identifier + " = " + offsets.GenerateCode() + ";";
			}
		}

		private class LocalReferenceExpression : Expression
		{
			private string identifier;

			public LocalReferenceExpression(string identifier)
			{
				this.identifier = identifier;
			}

			public override string GenerateCode()
			{
				return identifier;
			}
		}

		private class LocalAssignmentStatement : FunctionBodyStatement
		{
			private string identifier;
			private Expression expression;

			public LocalAssignmentStatement(string identifier, Expression expression)
			{
				this.identifier = identifier;
				this.expression = expression;
			}

			public override string GenerateCode()
			{
				return "double " + identifier + " = " + expression.GenerateCode() + ";";
			}
		}

		private class FunctionValueAssignmentStatement : FunctionBodyStatement
		{
			private string identifier;
			private int offset;
			private Expression expression;

			public FunctionValueAssignmentStatement(string identifier, int offset, Expression expression)
			{
				this.identifier = identifier;
				this.offset = offset;
				this.expression = expression;
			}

			public override string GenerateCode()
			{
				StringBuilder sb = new StringBuilder();
				sb.Append("a[");
				sb.Append(identifier);
				if (offset != 0)
				{
					sb.Append(" + ");
					sb.Append(offset);
				}
				sb.Append("] = ");
				sb.Append(expression.GenerateCode());
				sb.Append(";");

				return sb.ToString();
			}
		}

		private class LoopZeroStatement : FunctionBodyStatement
		{
			private string identifier;
			private int index, length;

			public LoopZeroStatement(string identifier, int index, int length)
			{
				this.identifier = identifier;
				this.index = index;
				this.length = length;
			}

			public override string GenerateCode()
			{
				return "for (int i = 0; i < " + length + "; i++)\r\n{\r\n    a[" + identifier + " + " + index + " + i] = 0.0;\r\n}";
			}
		}

		private class GradientLoopStatement : FunctionBodyStatement
		{
			private string identifier;
			private StorageOffsetsReference offsets;
			private GradientLoopTerm[] terms;
			private bool addition;

			public GradientLoopStatement(string identifier, StorageOffsetsReference offsets, GradientLoopTerm[] terms, bool addition)
			{
				this.identifier = identifier;
				this.offsets = offsets;
				this.terms = terms;
				this.addition = addition;
			}

			public override string GenerateCode()
			{
				StringBuilder sb = new StringBuilder();

				sb.Append("for (int i = 0; i < ");
				sb.Append(offsets.Count);
				sb.AppendLine("; i++)");
				sb.AppendLine("{");
				sb.Append("    a[");
				sb.Append(identifier);
				sb.Append(" + ");
				sb.Append(offsets.Identifier);
				sb.Append("[i]] ");
				if (addition)
				{
					sb.Append("+");
				}
				sb.Append("= ");

				for (int i = 0; i < terms.Length; i++)
				{
					if (i > 0)
					{
						sb.Append(" + ");
					}
					sb.Append(terms[i].GenerateCode());
				}

				sb.AppendLine(";");
				sb.AppendLine("}");

				return sb.ToString();
			}
		}

		private class HessianLoopStatement : FunctionBodyStatement
		{
			private string identifier;
			private StorageOffsetsReference offsets;
			private HessianLoopTerm[] terms;
			private bool addition;

			public HessianLoopStatement(string identifier, StorageOffsetsReference offsets, HessianLoopTerm[] terms, bool addition)
			{
				this.identifier = identifier;
				this.offsets = offsets;
				this.terms = terms;
				this.addition = addition;
			}

			public override string GenerateCode()
			{
				StringBuilder sb = new StringBuilder();

				sb.Append("for (int i = 0; i < ");
				sb.Append(offsets.Count);
				sb.AppendLine("; i++)");
				sb.AppendLine("{");
				sb.Append("    a[");
				sb.Append(identifier);
				sb.Append(" + ");
				sb.Append(offsets.Identifier);
				sb.Append("[i]] ");
				if (addition)
				{
					sb.Append("+");
				}
				sb.Append("= ");

				for (int i = 0; i < terms.Length; i++)
				{
					if (i > 0)
					{
						sb.Append(" + ");
					}
					sb.Append(terms[i].GenerateCode());
				}

				sb.AppendLine(";");
				sb.AppendLine("}");

				return sb.ToString();
			}
		}
	}
}
